[Feature] Add support for loading datasets from local Minari cache #3068

Ibinarriaga8 · 2025-07-14T11:13:34Z

Description

This PR adds support for loading datasets directly from the local Minari cache in the MinariExperienceReplay class. Specifically, it introduces the load_from_local_minari argument, which, when set to True, instructs the class to load the dataset from the user's local Minari cache (typically at ~/.minari/datasets) and skip any fetching from the Minari server (i.e., no remote download or overwrite will occur). After loading from the local cache, all subsequent preprocessing and loading steps continue as usual, ensuring the dataset is processed and made available correctly. This is especially useful for custom Minari datasets or when you want to avoid network access.

Documentation has been updated in the class docstring to clearly state the new behavior of the load_from_local_minari argument, including details about local cache prioritization and unchanged downstream preprocessing.

This PR also includes comprehensive test coverage for the new feature, confirming that datasets created and stored in the local Minari cache can be loaded, sampled, and validated for correctness using the new argument. The provided test (test_local_minari_dataset_loading) creates a custom dataset, loads it from cache, verifies sample integrity, and cleans up afterwards.

Motivation and Context

Previously, MinariExperienceReplay required datasets to be downloaded via its own interface, which was incompatible with custom or preloaded Minari datasets, attempting to load these would result in a FileNotFoundError.
This change allows users to work with their own datasets, datasets created with minari.DataCollector(...).create_dataset(...), or any dataset present in the local Minari cache, without requiring redundant downloads or manual metadata copying.

If the dataset is not found in the local cache, a FileNotFoundError is raised with a clear message.

I have raised an issue to propose this change (required for new features and bug fixes)

Solves #3067

Types of changes

New feature (non-breaking change which adds core functionality)
Documentation (update in the documentation)

Checklist

I have read the CONTRIBUTION guide (required)
My change requires a change to the documentation.
I have updated the tests accordingly (required for a bug fix or a new feature).
I have updated the documentation accordingly.

pytorch-bot · 2025-07-14T11:13:38Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3068

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Ibinarriaga8 · 2025-07-14T11:27:36Z

test/test_libs.py

+
+        dataset_id = "cartpole/test-local-v1"
+
+        # Create dataset using Gym + DataCollector


Custom minari dataset creation from a gymnasium environment

vmoens

That's really handy I love it!

vmoens · 2025-07-14T13:09:17Z

@Ibinarriaga8

Do you think we have an opportunity here to reduce the number of datasets we download to test minari, and use customly built datasets instead like you do in your test?

Ibinarriaga8 · 2025-07-14T13:40:48Z

@Ibinarriaga8

Do you think we have an opportunity here to reduce the number of datasets we download to test minari, and use customly built datasets instead like you do in your test?

Absolutely, instead of downloading 20 datasets from the Minari server, we could generate smaller, custom datasets from any gymnasium environment as part of our tests. This approach is especially valuable for D4RL datasets, which tend to be very large and can significantly slow down testing.

vmoens · 2025-07-14T13:52:43Z

test/test_libs.py

@@ -29,9 +29,12 @@
 from sys import platform
 from unittest import mock

+import minari


This is a global import - we must avoid these at all cost.
Can you make it local?

Ok, I can fix it

vmoens · 2025-07-14T13:53:21Z

Absolutely, instead of downloading 20 datasets from the Minari server, we could generate smaller, custom datasets from any gymnasium environment as part of our tests. This approach is especially valuable for D4RL datasets, which tend to be very large and can significantly slow down testing.

Do you want to give it a go? Otherwise I can do it no worry

Ibinarriaga8 · 2025-07-14T13:55:10Z

Absolutely, instead of downloading 20 datasets from the Minari server, we could generate smaller, custom datasets from any gymnasium environment as part of our tests. This approach is especially valuable for D4RL datasets, which tend to be very large and can significantly slow down testing.

Do you want to give it a go? Otherwise I can do it no worry

Yes, no problem

Ibinarriaga8 · 2025-07-14T14:26:41Z

@vmoens Do you think we should replace testing with all Minari datasets by using smaller, custom datasets generated from gymnasium environments? Or should we still download and test with some datasets from the Minari server to ensure that the downloading functionality in MinariExperienceReplay works correctly?

vmoens · 2025-07-14T15:25:48Z

@vmoens Do you think we should replace testing with all Minari datasets by using smaller, custom datasets generated from gymnasium environments? Or should we still download and test with some datasets from the Minari server to ensure that the downloading functionality in MinariExperienceReplay works correctly?

Maybe just one defined (not random) dataset - to make sure?

test/test_libs.py

Ibinarriaga8 · 2025-07-14T16:50:33Z

@vmoens Do you think we should replace testing with all Minari datasets by using smaller, custom datasets generated from gymnasium environments? Or should we still download and test with some datasets from the Minari server to ensure that the downloading functionality in MinariExperienceReplay works correctly?

Maybe just one defined (not random) dataset - to make sure?

Ok, I agree that testing with a single, well-defined dataset makes sense to ensure the download and loading functionality are covered.

However, with most gym environments, creating custom datasets isn’t always straightforward. For example, I ran into this error:

ValueError: Dict key mismatch; expected keys: ['reward_ctrl', 'reward_forward', 'reward_survive', 'x_position', 'x_velocity', 'z_distance_from_origin']; dict: {'x_position': np.float64(0.001823518632481435), 'z_distance_from_origin': np.float64(-0.004461789811977868)}

I’ll keep investigating to see if there’s a workaround or a more general approach that works across environments.

vmoens · 2025-07-15T12:43:35Z

Ok maybe we could proceed with this in its current state and refactor the tests later!

Ibinarriaga8 · 2025-07-15T13:33:53Z

Ok maybe we could proceed with this in its current state and refactor the tests later!

Ok, I have refactored the tests to use this custom datasets approach. I am still fixing some minor errors and making sure it all works.
I can either commit these changes to this PR or create a separate PR for the test refactor. Whichever you prefer

vmoens · 2025-07-15T15:30:30Z

If you're ready feel free to commit here directly

vmoens

LGTM!
This actually makes me think that we should have a similar pipeline where one creates a datacollector with some replay buffer and serializes it as a dataset later on.
It's mostly a matter of documenting it rather than creating the feature, everything already exists I think!

Ibinarriaga8 · 2025-07-16T13:02:11Z

LGTM! This actually makes me think that we should have a similar pipeline where one creates a datacollector with some replay buffer and serializes it as a dataset later on. It's mostly a matter of documenting it rather than creating the feature, everything already exists I think!

Ok, should I provide a tutorial documenting the pipeline from DataCollector with a replay buffer to dataset serialization?

vmoens · 2025-07-17T00:13:18Z

First let's fix this!
From the CI I see there are some global imports of ale_py etc occurring in every workflow. Probably because Minari is imported in some function that is called in the global workspace.
Every Minari import should be localised within the tests

Ibinarriaga8 · 2025-07-17T09:46:44Z

First let's fix this! From the CI I see there are some global imports of ale_py etc occurring in every workflow. Probably because Minari is imported in some function that is called in the global workspace. Every Minari import should be localised within the tests

I already import minari and ale_py only inside the specific functions or test cases that require them, rather than at the global level.
Is there a particular place where you’re seeing these imports happen globally, or could it be that the function importing them is being executed in the global workspace?

Ibinarriaga8 · 2025-07-17T09:50:41Z

First let's fix this! From the CI I see there are some global imports of ale_py etc occurring in every workflow. Probably because Minari is imported in some function that is called in the global workspace. Every Minari import should be localised within the tests

I already import minari and ale_py only inside the specific functions or test cases that require them, rather than at the global level. Is there a particular place where you’re seeing these imports happen globally, or could it be that the function importing them is being executed in the global workspace?

Do you think moving all imports to the _minari_init which only makes necessary imports if dependencies are installed would solve this?

vmoens · 2025-07-17T15:18:06Z

No imports must happen if and only if they're needed. Otherwise spawned multiprocessed jobs all import a ton of useless libs (if they're installed).
We only want local imports without try except blocks.

Ibinarriaga8 · 2025-07-15T15:53:27Z

test/test_libs.py

@@ -3341,6 +3341,39 @@ def test_d4rl_iteration(self, task, split_trajs):

 _MINARI_DATASETS = []

+MUJOCO_ENVIRONMENTS = [


Explicitly specified current set of Minari supported datasets for integration with Gym environments

Specified datasets are those used in Mujoco and D4RL

Ibinarriaga8 · 2025-07-15T15:56:03Z

test/test_libs.py

-# Initialize with placeholder values for parametrization
-# These will be replaced with actual dataset names when the first Minari test runs
-_MINARI_DATASETS = [str(i) for i in range(20)]
+def get_random_minigrid_datasets():


Only use Minari datasets from Minigrid

This is because current version of Minari cannot serialize custom MissionSpace objects, which are used in most Minigrid environments

This means you cannot create your custom dataset from a minigrid environment directly; you will need to modify Mission space

Ibinarriaga8 · 2025-07-15T15:57:56Z

test/test_libs.py

+            cleanup_needed = False
+
+        else:
+            # Atari environment datasets


For Atari environments, select a random subset of environments.

Ibinarriaga8 · 2025-07-15T16:02:28Z

test/test_libs.py

@@ -29,11 +29,9 @@
 from sys import platform
 from unittest import mock

-


Removed global imports

Ibinarriaga8 · 2025-07-15T16:03:14Z

test/test_libs.py

Note: new dependencies were added for creating custom Minari datasets from gym environments:

ale_py (for Atari)

gymnasium_robotics (for D4RL & Mujoco)

Ibinarriaga8 · 2025-07-17T15:45:03Z

test/test_libs.py

+        range(
+            len(MUJOCO_ENVIRONMENTS)
+            + len(D4RL_ENVIRONMENTS)
+            + len(get_random_minigrid_datasets())


These two functions import minari and are called at the module level (global workspace), which causes the error during CI.
I will refactor them so that the imports and calls only occur within the relevant test functions.

Ibinarriaga8 · 2025-07-18T13:51:24Z

test/test_libs.py

@@ -3488,12 +3488,7 @@ class TestMinari:
    @pytest.mark.parametrize("split", [False, True])
    @pytest.mark.parametrize(
        "dataset_idx",
-        range(
-            len(MUJOCO_ENVIRONMENTS)


These decorators were called in global workspace and imported minari and ale_py internally

Now, I use a static upper bound, so that no function that imports minari or ale_py globally is called

vmoens · 2025-07-21T15:43:56Z

Cool! Looks like the conflicts with other CI runs are resolved. We have some long-lasting failures but I wouldn't pay too much attention to them.
The minari tests are failing though - do you want me to help? I see a "too many requests" error which I think can only be solved by having fewer tests and a missing dependency to be added here

Ibinarriaga8 · 2025-07-21T21:53:51Z

test/test_libs.py

Reduce number of fetched datasets to avoid 429 Error

Ibinarriaga8 · 2025-07-21T21:54:41Z

.github/unittest/linux_libs/scripts_minari/environment.yml

@@ -21,3 +21,5 @@ dependencies:
    - hydra-core
    - minari[gcs,hdf5,hf]
    - gymnasium<1.0.0
+    - ale-py


Added new dependencies

Ibinarriaga8 added 2 commits July 14, 2025 11:29

[Feature] Add support for loading datasets from local Minari cache

26aa0ff

[Refactor] Fixed linting errors

a56c508

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 14, 2025

Ibinarriaga8 commented Jul 14, 2025

View reviewed changes

vmoens added the enhancement New feature or request label Jul 14, 2025

vmoens approved these changes Jul 14, 2025

View reviewed changes

vmoens added Environments Adds or modifies an environment wrapper Data Data-related PR, will launch data-related jobs labels Jul 14, 2025

vmoens reviewed Jul 14, 2025

View reviewed changes

[Refactor] moved minari import to local scope

e90f4d0

vmoens reviewed Jul 14, 2025

View reviewed changes

test/test_libs.py Outdated Show resolved Hide resolved

[Refactor] Refactor Minari testing to use custom datasets

cc43f9a

vmoens approved these changes Jul 16, 2025

View reviewed changes

[Fix] eliminate global imports for test_libs

23aecec

Ibinarriaga8 commented Jul 18, 2025

View reviewed changes

Ibinarriaga8 and others added 3 commits July 18, 2025 16:00

[Refactor] Reduce dataset index range in Minari tests

9321adb

Merge branch 'pytorch:main' into main

a4993c3

Merge branch 'pytorch:main' into main

08779c2

Ibinarriaga8 added 2 commits July 21, 2025 23:51

[Enhancement] Update minari environment.yml and refactored minari test

0ab8f81

[Refactor] Reduced datasets for minari tests

01965f4

Ibinarriaga8 commented Jul 21, 2025

View reviewed changes

test/test_libs.py

Copy link

Contributor Author

Ibinarriaga8 Jul 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reduce number of fetched datasets to avoid 429 Error

Ibinarriaga8 commented Jul 21, 2025

View reviewed changes


		dataset_id = "cartpole/test-local-v1"

		# Create dataset using Gym + DataCollector

		@@ -3341,6 +3341,39 @@ def test_d4rl_iteration(self, task, split_trajs):

		_MINARI_DATASETS = []

		MUJOCO_ENVIRONMENTS = [

		@@ -29,11 +29,9 @@
		from sys import platform
		from unittest import mock

[Feature] Add support for loading datasets from local Minari cache #3068

Are you sure you want to change the base?

[Feature] Add support for loading datasets from local Minari cache #3068

Conversation

Ibinarriaga8 commented Jul 14, 2025

Description

Motivation and Context

Types of changes

Checklist

Uh oh!

pytorch-bot bot commented Jul 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3068

Uh oh!

Ibinarriaga8 Jul 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

vmoens commented Jul 14, 2025

Uh oh!

Ibinarriaga8 commented Jul 14, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vmoens commented Jul 14, 2025

Uh oh!

Ibinarriaga8 commented Jul 14, 2025

Uh oh!

Ibinarriaga8 commented Jul 14, 2025

Uh oh!

vmoens commented Jul 14, 2025

Uh oh!

Uh oh!

Ibinarriaga8 commented Jul 14, 2025

Uh oh!

vmoens commented Jul 15, 2025

Uh oh!

Ibinarriaga8 commented Jul 15, 2025

Uh oh!

vmoens commented Jul 15, 2025

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

Ibinarriaga8 commented Jul 16, 2025

Uh oh!

vmoens commented Jul 17, 2025

Uh oh!

Ibinarriaga8 commented Jul 17, 2025

Uh oh!

Ibinarriaga8 commented Jul 17, 2025

Uh oh!

vmoens commented Jul 17, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vmoens commented Jul 21, 2025

Uh oh!

pytorch-bot bot commented Jul 14, 2025 •

edited

Loading

Ibinarriaga8 Jul 14, 2025 •

edited

Loading